//! 负载均衡指标收集模块
//! 
//! 提供性能指标收集和监控功能

use std::sync::Arc;
use std::time::{Duration, Instant};
use std::collections::HashMap;
use tokio::sync::RwLock;
use anyhow::Result;
use tracing::{info, debug};

/// 负载均衡器指标
pub struct LoadBalancerMetrics {
    /// 总请求数
    total_requests: Arc<RwLock<u64>>,
    /// 成功请求数
    successful_requests: Arc<RwLock<u64>>,
    /// 失败请求数
    failed_requests: Arc<RwLock<u64>>,
    /// 错误计数（按类型）
    error_counts: Arc<RwLock<HashMap<String, u64>>>,
    /// 后端请求计数
    backend_requests: Arc<RwLock<HashMap<String, BackendMetrics>>>,
    /// 选择时间统计
    selection_times: Arc<RwLock<Vec<Duration>>>,
    /// 启动时间
    start_time: Instant,
    /// 最后重置时间
    last_reset: Arc<RwLock<Instant>>,
}

/// 后端指标
#[derive(Debug, Clone)]
pub struct BackendMetrics {
    /// 请求数
    pub request_count: u64,
    /// 成功请求数
    pub success_count: u64,
    /// 失败请求数
    pub error_count: u64,
    /// 平均响应时间
    pub avg_response_time: Duration,
    /// 最大响应时间
    pub max_response_time: Duration,
    /// 最小响应时间
    pub min_response_time: Duration,
    /// 响应时间样本
    response_times: Vec<Duration>,
}

/// 指标快照
#[derive(Debug, Clone)]
pub struct MetricsSnapshot {
    /// 总请求数
    pub total_requests: u64,
    /// 成功请求数
    pub successful_requests: u64,
    /// 失败请求数
    pub failed_requests: u64,
    /// 成功率
    pub success_rate: f64,
    /// 错误计数
    pub error_counts: HashMap<String, u64>,
    /// 后端指标
    pub backend_metrics: HashMap<String, BackendMetrics>,
    /// 平均选择时间
    pub avg_selection_time: Duration,
    /// 运行时间
    pub uptime: Duration,
    /// 请求率（每秒）
    pub requests_per_second: f64,
}

impl BackendMetrics {
    /// 创建新的后端指标
    pub fn new() -> Self {
        Self {
            request_count: 0,
            success_count: 0,
            error_count: 0,
            avg_response_time: Duration::from_millis(0),
            max_response_time: Duration::from_millis(0),
            min_response_time: Duration::from_secs(u64::MAX),
            response_times: Vec::new(),
        }
    }

    /// 记录请求
    pub fn record_request(&mut self, success: bool, response_time: Option<Duration>) {
        self.request_count += 1;
        
        if success {
            self.success_count += 1;
        } else {
            self.error_count += 1;
        }

        if let Some(time) = response_time {
            self.response_times.push(time);
            
            // 更新最大最小值
            if time > self.max_response_time {
                self.max_response_time = time;
            }
            if time < self.min_response_time {
                self.min_response_time = time;
            }

            // 计算平均响应时间
            if !self.response_times.is_empty() {
                let total: Duration = self.response_times.iter().sum();
                self.avg_response_time = total / self.response_times.len() as u32;
            }

            // 限制样本数量，避免内存过度使用
            if self.response_times.len() > 1000 {
                self.response_times.drain(0..500); // 移除前一半样本
            }
        }
    }

    /// 获取成功率
    pub fn success_rate(&self) -> f64 {
        if self.request_count == 0 {
            1.0
        } else {
            self.success_count as f64 / self.request_count as f64
        }
    }
}

impl Default for BackendMetrics {
    fn default() -> Self {
        Self::new()
    }
}

impl LoadBalancerMetrics {
    /// 创建新的指标收集器
    pub fn new() -> Self {
        Self {
            total_requests: Arc::new(RwLock::new(0)),
            successful_requests: Arc::new(RwLock::new(0)),
            failed_requests: Arc::new(RwLock::new(0)),
            error_counts: Arc::new(RwLock::new(HashMap::new())),
            backend_requests: Arc::new(RwLock::new(HashMap::new())),
            selection_times: Arc::new(RwLock::new(Vec::new())),
            start_time: Instant::now(),
            last_reset: Arc::new(RwLock::new(Instant::now())),
        }
    }

    /// 启动指标收集
    pub async fn start(&self) -> Result<()> {
        info!("启动负载均衡指标收集");
        Ok(())
    }

    /// 停止指标收集
    pub async fn stop(&self) -> Result<()> {
        info!("停止负载均衡指标收集");
        Ok(())
    }

    /// 增加请求计数
    pub async fn increment_request_count(&self) {
        let mut total = self.total_requests.write().await;
        *total += 1;
        debug!("请求计数增加，当前总数: {}", *total);
    }

    /// 增加成功请求计数
    pub async fn increment_success_count(&self) {
        let mut successful = self.successful_requests.write().await;
        *successful += 1;
    }

    /// 增加失败请求计数
    pub async fn increment_failure_count(&self) {
        let mut failed = self.failed_requests.write().await;
        *failed += 1;
    }

    /// 增加错误计数
    pub async fn increment_error_count(&self, error_type: &str) {
        let mut errors = self.error_counts.write().await;
        *errors.entry(error_type.to_string()).or_insert(0) += 1;
        
        // 同时增加失败请求计数
        self.increment_failure_count().await;
    }

    /// 记录后端请求
    pub async fn increment_backend_requests(&self, backend_id: &str) {
        let mut backend_metrics = self.backend_requests.write().await;
        let metrics = backend_metrics.entry(backend_id.to_string()).or_insert_with(BackendMetrics::new);
        metrics.record_request(true, None);
    }

    /// 记录后端请求结果
    pub async fn record_backend_request(&self, backend_id: &str, success: bool, response_time: Option<Duration>) {
        let mut backend_metrics = self.backend_requests.write().await;
        let metrics = backend_metrics.entry(backend_id.to_string()).or_insert_with(BackendMetrics::new);
        metrics.record_request(success, response_time);

        if success {
            self.increment_success_count().await;
        } else {
            self.increment_failure_count().await;
        }
    }

    /// 记录选择时间
    pub async fn record_selection_time(&self, duration: Duration) {
        let mut times = self.selection_times.write().await;
        times.push(duration);

        // 限制样本数量
        if times.len() > 1000 {
            times.drain(0..500);
        }
    }

    /// 获取总请求数
    pub async fn get_total_requests(&self) -> u64 {
        *self.total_requests.read().await
    }

    /// 获取成功请求数
    pub async fn get_successful_requests(&self) -> u64 {
        *self.successful_requests.read().await
    }

    /// 获取失败请求数
    pub async fn get_failed_requests(&self) -> u64 {
        *self.failed_requests.read().await
    }

    /// 获取成功率
    pub async fn get_success_rate(&self) -> f64 {
        let total = self.get_total_requests().await;
        let successful = self.get_successful_requests().await;

        if total == 0 {
            1.0
        } else {
            successful as f64 / total as f64
        }
    }

    /// 获取平均选择时间
    pub async fn get_avg_selection_time(&self) -> Duration {
        let times = self.selection_times.read().await;
        if times.is_empty() {
            Duration::from_millis(0)
        } else {
            let total: Duration = times.iter().sum();
            total / times.len() as u32
        }
    }

    /// 获取请求率（每秒）
    pub async fn get_requests_per_second(&self) -> f64 {
        let total_requests = self.get_total_requests().await;
        let uptime = self.start_time.elapsed();
        
        if uptime.as_secs() == 0 {
            0.0
        } else {
            total_requests as f64 / uptime.as_secs() as f64
        }
    }

    /// 获取指标快照
    pub async fn get_snapshot(&self) -> MetricsSnapshot {
        let total_requests = self.get_total_requests().await;
        let successful_requests = self.get_successful_requests().await;
        let failed_requests = self.get_failed_requests().await;
        let success_rate = self.get_success_rate().await;
        let avg_selection_time = self.get_avg_selection_time().await;
        let requests_per_second = self.get_requests_per_second().await;
        
        let error_counts = self.error_counts.read().await.clone();
        let backend_metrics = self.backend_requests.read().await.clone();
        
        MetricsSnapshot {
            total_requests,
            successful_requests,
            failed_requests,
            success_rate,
            error_counts,
            backend_metrics,
            avg_selection_time,
            uptime: self.start_time.elapsed(),
            requests_per_second,
        }
    }

    /// 重置指标
    pub async fn reset(&self) {
        *self.total_requests.write().await = 0;
        *self.successful_requests.write().await = 0;
        *self.failed_requests.write().await = 0;
        self.error_counts.write().await.clear();
        self.backend_requests.write().await.clear();
        self.selection_times.write().await.clear();
        *self.last_reset.write().await = Instant::now();
        
        info!("负载均衡指标已重置");
    }

    /// 获取后端指标
    pub async fn get_backend_metrics(&self, backend_id: &str) -> Option<BackendMetrics> {
        let backend_metrics = self.backend_requests.read().await;
        backend_metrics.get(backend_id).cloned()
    }

    /// 获取所有后端指标
    pub async fn get_all_backend_metrics(&self) -> HashMap<String, BackendMetrics> {
        self.backend_requests.read().await.clone()
    }

    /// 移除后端指标
    pub async fn remove_backend_metrics(&self, backend_id: &str) {
        let mut backend_metrics = self.backend_requests.write().await;
        backend_metrics.remove(backend_id);
    }

    /// 导出Prometheus格式指标
    pub async fn export_prometheus_metrics(&self) -> String {
        let snapshot = self.get_snapshot().await;
        let mut output = String::new();

        // 基础指标
        output.push_str(&format!(
            "# HELP lb_requests_total Total number of requests\n\
             # TYPE lb_requests_total counter\n\
             lb_requests_total {}\n\n",
            snapshot.total_requests
        ));

        output.push_str(&format!(
            "# HELP lb_requests_successful_total Total number of successful requests\n\
             # TYPE lb_requests_successful_total counter\n\
             lb_requests_successful_total {}\n\n",
            snapshot.successful_requests
        ));

        output.push_str(&format!(
            "# HELP lb_requests_failed_total Total number of failed requests\n\
             # TYPE lb_requests_failed_total counter\n\
             lb_requests_failed_total {}\n\n",
            snapshot.failed_requests
        ));

        output.push_str(&format!(
            "# HELP lb_success_rate Success rate of requests\n\
             # TYPE lb_success_rate gauge\n\
             lb_success_rate {:.4}\n\n",
            snapshot.success_rate
        ));

        output.push_str(&format!(
            "# HELP lb_selection_time_avg_seconds Average backend selection time\n\
             # TYPE lb_selection_time_avg_seconds gauge\n\
             lb_selection_time_avg_seconds {:.6}\n\n",
            snapshot.avg_selection_time.as_secs_f64()
        ));

        output.push_str(&format!(
            "# HELP lb_requests_per_second Request rate per second\n\
             # TYPE lb_requests_per_second gauge\n\
             lb_requests_per_second {:.2}\n\n",
            snapshot.requests_per_second
        ));

        // 错误指标
        for (error_type, count) in &snapshot.error_counts {
            output.push_str(&format!(
                "# HELP lb_errors_total Total number of errors by type\n\
                 # TYPE lb_errors_total counter\n\
                 lb_errors_total{{error_type=\"{}\"}} {}\n\n",
                error_type, count
            ));
        }

        // 后端指标
        for (backend_id, metrics) in &snapshot.backend_metrics {
            output.push_str(&format!(
                "# HELP lb_backend_requests_total Total requests per backend\n\
                 # TYPE lb_backend_requests_total counter\n\
                 lb_backend_requests_total{{backend=\"{}\"}} {}\n\n",
                backend_id, metrics.request_count
            ));

            output.push_str(&format!(
                "# HELP lb_backend_success_rate Success rate per backend\n\
                 # TYPE lb_backend_success_rate gauge\n\
                 lb_backend_success_rate{{backend=\"{}\"}} {:.4}\n\n",
                backend_id, metrics.success_rate()
            ));

            output.push_str(&format!(
                "# HELP lb_backend_response_time_avg_seconds Average response time per backend\n\
                 # TYPE lb_backend_response_time_avg_seconds gauge\n\
                 lb_backend_response_time_avg_seconds{{backend=\"{}\"}} {:.6}\n\n",
                backend_id, metrics.avg_response_time.as_secs_f64()
            ));
        }

        output
    }
}

impl Default for LoadBalancerMetrics {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn test_metrics_basic_operations() {
        let metrics = LoadBalancerMetrics::new();

        // 测试请求计数
        metrics.increment_request_count().await;
        metrics.increment_request_count().await;
        
        assert_eq!(metrics.get_total_requests().await, 2);

        // 测试成功/失败计数
        metrics.increment_success_count().await;
        metrics.increment_failure_count().await;

        assert_eq!(metrics.get_successful_requests().await, 1);
        assert_eq!(metrics.get_failed_requests().await, 1);
        assert_eq!(metrics.get_success_rate().await, 0.5);
    }

    #[tokio::test]
    async fn test_backend_metrics() {
        let metrics = LoadBalancerMetrics::new();
        
        // 记录后端请求
        metrics.record_backend_request(
            "backend1",
            true,
            Some(Duration::from_millis(100))
        ).await;
        
        metrics.record_backend_request(
            "backend1",
            false,
            Some(Duration::from_millis(200))
        ).await;

        let backend_metrics = metrics.get_backend_metrics("backend1").await;
        assert!(backend_metrics.is_some());
        
        let backend_metrics = backend_metrics.unwrap();
        assert_eq!(backend_metrics.request_count, 2);
        assert_eq!(backend_metrics.success_count, 1);
        assert_eq!(backend_metrics.error_count, 1);
        assert_eq!(backend_metrics.success_rate(), 0.5);
    }

    #[tokio::test]
    async fn test_selection_time_recording() {
        let metrics = LoadBalancerMetrics::new();
        
        metrics.record_selection_time(Duration::from_millis(10)).await;
        metrics.record_selection_time(Duration::from_millis(20)).await;
        
        let avg_time = metrics.get_avg_selection_time().await;
        assert_eq!(avg_time, Duration::from_millis(15));
    }

    #[tokio::test]
    async fn test_metrics_snapshot() {
        let metrics = LoadBalancerMetrics::new();
        
        // 添加一些数据
        metrics.increment_request_count().await;
        metrics.increment_success_count().await;
        metrics.increment_error_count("timeout").await;
        
        let snapshot = metrics.get_snapshot().await;
        
        assert_eq!(snapshot.total_requests, 1);
        assert_eq!(snapshot.successful_requests, 1);
        assert_eq!(snapshot.failed_requests, 1);
        assert_eq!(snapshot.error_counts.get("timeout"), Some(&1));
    }

    #[tokio::test]
    async fn test_metrics_reset() {
        let metrics = LoadBalancerMetrics::new();
        
        // 添加数据
        metrics.increment_request_count().await;
        metrics.increment_success_count().await;
        
        assert_eq!(metrics.get_total_requests().await, 1);
        
        // 重置
        metrics.reset().await;
        
        assert_eq!(metrics.get_total_requests().await, 0);
        assert_eq!(metrics.get_successful_requests().await, 0);
    }

    #[test]
    fn test_prometheus_export() {
        tokio_test::block_on(async {
            let metrics = LoadBalancerMetrics::new();
            
            metrics.increment_request_count().await;
            metrics.increment_success_count().await;
            
            let prometheus_output = metrics.export_prometheus_metrics().await;
            
            assert!(prometheus_output.contains("lb_requests_total 1"));
            assert!(prometheus_output.contains("lb_requests_successful_total 1"));
        });
    }
}